package com.wrt.produce.source;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import javax.sql.DataSource;

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

/**
 * 动态数据源
 * 
 * @author 文瑞涛
 * @date 2021年9月6日 下午10:26:58
 */
public class DynamicDataSource extends AbstractRoutingDataSource {

	private static final String DEFAULT_SOURCE_KEY = "defaultDataSource";

	private Map<Object, Object> dataSourcesMap = new ConcurrentHashMap<>(10);

	private final ThreadLocal<String> contextHolder = ThreadLocal.withInitial(() -> DEFAULT_SOURCE_KEY);

	public DynamicDataSource(DataSource defaultDataSource) {
		dataSourcesMap.put(DEFAULT_SOURCE_KEY, defaultDataSource);
		setTargetDataSources(dataSourcesMap);
		setDefaultTargetDataSource(defaultDataSource);
	}

	@Override
	protected Object determineCurrentLookupKey() {
		return contextHolder.get();
	}

	public void addDataSource(String key, DataSource datasource) {
		dataSourcesMap.put(key, datasource);
		setDataSourceKey(key);
	}

	public Map<Object, Object> getDataSourcesMap() {
		return this.dataSourcesMap;
	}

	public DataSource getDataSource(Object key) {
		return (DataSource) this.dataSourcesMap.get(key);
	}

	public DataSource getDataSource() {
		return (DataSource) this.dataSourcesMap.get(determineCurrentLookupKey());
	}

	/**
	 * 切换数据源
	 * 
	 * @param key
	 */
	public void setDataSourceKey(String key) {
		contextHolder.set(key);
		this.afterPropertiesSet();
	}

	public void removeDataSource(Object key) {
		if (this.dataSourcesMap.containsKey(key)) {
			this.dataSourcesMap.remove(key);
		}
	}

	/**
	 * 切换到默认数据源
	 */
	public void clearDataSourceKey() {
		contextHolder.remove();
	}
}
